{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "74bc2ccb",
   "metadata": {},
   "source": [
    "# Explore large language models on Apple silicon with MLX"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9f4b67f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"TOKENIZERS_PARALLELISM\"]=\"false\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d23eda5f",
   "metadata": {},
   "source": [
    "### Demo 1: Running DeepSeek AI’s latest model with 670 billion parameters.\n",
    "* Note 1: This example requires Mac Studio M3 Ultra with 512 GB of unified memory.\n",
    "* Note 2: Copy paste the line below and run it in the terminal, since Jupyter Notebook output doesn't allow turn-by-turn chat iteraction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a45bf5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run this command in the terminal to chat with `DeepSeek-V3-0324-4bit`\n",
    "#mlx_lm.chat --model mlx-community/DeepSeek-V3-0324-4bit"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "292a968d",
   "metadata": {},
   "source": [
    "### Using the `mlx_lm.generate` command"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "38f6b4a1",
   "metadata": {},
   "source": [
    "Easiest way to generate text with LLMs is to use the `mlx_lm.generate` command"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51bd2ed4",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mlx_lm.generate --model \"mlx-community/Mistral-7B-Instruct-v0.3-4bit\" \\\n",
    "                 --prompt \"Write a quick sort in Swift\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1fa6be17",
   "metadata": {},
   "source": [
    "You can tweak the behavior of the model by adding flags for things like sampling temperature, top-p, or max tokens; just like with any standard text generation setup."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f7add212",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mlx_lm.generate --model \"mlx-community/Mistral-7B-Instruct-v0.3-4bit\" \\\n",
    "                 --prompt \"Write a quick sort in Swift\" \\\n",
    "                 --top-p 0.5 \\\n",
    "                 --temp 0.2 \\\n",
    "                 --max-tokens 1024"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "761bb876",
   "metadata": {},
   "source": [
    " ### Using the Python API\n",
    " Using the flexible Python API for fine-grained control and to integrate generation into a larger workflows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e042a321",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Using MLX LM from Python\n",
    "\n",
    "from mlx_lm import load, generate\n",
    "\n",
    "# Load the model and tokenizer directly from HF\n",
    "model, tokenizer = load(\"mlx-community/Mistral-7B-Instruct-v0.3-4bit\")\n",
    "\n",
    "# Prepare the prompt for the model\n",
    "prompt = \"Write a quick sort in Swift\"\n",
    "messages = [{\"role\": \"user\", \"content\": prompt}]\n",
    "prompt = tokenizer.apply_chat_template(\n",
    "    messages, add_generation_prompt=True\n",
    ")\n",
    "\n",
    "# Generate the text\n",
    "text = generate(model, tokenizer, prompt=prompt, verbose=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe5c4283",
   "metadata": {},
   "source": [
    "### Inspecting an mlx_lm model and exploring its architecture"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "629dfa50",
   "metadata": {},
   "outputs": [],
   "source": [
    "from mlx_lm import load, generate\n",
    "\n",
    "model, tokenizer = load(\"mlx-community/Mistral-7B-Instruct-v0.3-4bit\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7e7d43c",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "796d1c7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(model.parameters())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3b56bdc",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(model.layers[0].self_attn)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "64f8e6f2",
   "metadata": {},
   "source": [
    "### Generation with KV cache"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "775fd3f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from mlx_lm import load, generate\n",
    "from mlx_lm.models.cache import make_prompt_cache\n",
    "\n",
    "# Load the model and tokenizer directly from HF\n",
    "model, tokenizer = load(\"mlx-community/Mistral-7B-Instruct-v0.3-4bit\")\n",
    "\n",
    "# Prepare the prompt for the model\n",
    "prompt = \"Write a quick sort in Swift\"\n",
    "messages = [{\"role\": \"user\", \"content\": prompt}]\n",
    "prompt = tokenizer.apply_chat_template(\n",
    "    messages, add_generation_prompt=True\n",
    ")\n",
    "\n",
    "cache = make_prompt_cache(model)\n",
    "\n",
    "# Generate the text\n",
    "text = generate(model, tokenizer, prompt=prompt, prompt_cache=cache, verbose=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34d960f7",
   "metadata": {},
   "source": [
    "#### Let's ask a follow up question that the model can respond to using context in the KV cache"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d669073",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompt = \"how can I explain it to a five year old?\"\n",
    "messages = [{\"role\": \"user\", \"content\": prompt}]\n",
    "prompt = tokenizer.apply_chat_template(\n",
    "    messages, add_generation_prompt=True\n",
    ")\n",
    "\n",
    "# Use the cache to maintain context\n",
    "text = generate(model, tokenizer, prompt=prompt, prompt_cache=cache, verbose=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "569e3270",
   "metadata": {},
   "source": [
    "### Model quantization\n",
    "So far we've been using the 4-bit quantized version of the `Mistral-7b-Instruct-v0.3` model directly from the mlx-community on Hugging Face.\n",
    "Now we'll see how you can quantize the model using the `mlx_lm.convert` command.\n",
    "This tool takes care of downloading a model from Hugging Face, converting it to a different precision, and saving it locally — all in one step. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8218994",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "mlx_path=\"./mistral-7b-v0.3-4bit\"\n",
    "if not os.path.exists(mlx_path):\n",
    "    !mlx_lm.convert --hf-path \"mlx-community/Mistral-7B-Instruct-v0.3\" \\\n",
    "                --mlx-path \"./mistral-7b-v0.3-4bit\" \\\n",
    "                --dtype float16 \\\n",
    "                --quantize --q-bits 4 --q-group-size 64"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4e62b96",
   "metadata": {},
   "outputs": [],
   "source": [
    "import subprocess\n",
    "\n",
    "def get_directory_size_mb(path):\n",
    "    result = subprocess.run(['du', '-sm', path], stdout=subprocess.PIPE, text=True)\n",
    "    size_mb = int(result.stdout.split()[0])\n",
    "    size_gb = size_mb / 1024\n",
    "    return size_gb\n",
    "\n",
    "\n",
    "directory_path = os.path.expanduser('~/.cache/huggingface/hub/models--mlx-community--Mistral-7B-Instruct-v0.3')\n",
    "print(\"Size of original bfloat16 model\")\n",
    "print(\"===============================\")\n",
    "print(f\"{get_directory_size_mb(directory_path):2.4f} GB\")\n",
    "print()\n",
    "directory_path = './mistral-7b-v0.3-4bit'\n",
    "print(\"Size of quantized model\")\n",
    "print(\"===============================\")\n",
    "print(f\"{get_directory_size_mb(directory_path):2.4f} GB\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "944fe552",
   "metadata": {},
   "source": [
    "#### Apply different quantization settings to different parts of the model, all from Python"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9d2cd325",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Model quantization with MLX LM in Python\n",
    "\n",
    "from mlx_lm.convert import convert\n",
    "\n",
    "# We can choose a different quantization per layer\n",
    "def mixed_quantization(layer_path, layer, model_config):\n",
    "    if \"lm_head\" in layer_path or \"embed_tokens\" in layer_path:\n",
    "        return {\"bits\": 6, \"group_size\": 64}\n",
    "    elif hasattr(layer, \"to_quantized\"):\n",
    "        return {\"bits\": 4, \"group_size\": 64}\n",
    "    else:\n",
    "        return False\n",
    "\n",
    "# Convert can be used to change precision, quantize and upload models to HF\n",
    "mlx_path=\"./mistral-7b-v0.3-mixed-4-6-bit\"\n",
    "if not os.path.exists(mlx_path):\n",
    "    convert(\n",
    "        hf_path=\"mistralai/Mistral-7B-Instruct-v0.3\",\n",
    "        mlx_path=\"./mistral-7b-v0.3-mixed-4-6-bit\",\n",
    "        quantize=True,\n",
    "        quant_predicate=mixed_quantization\n",
    "    )\n",
    "\n",
    "print()\n",
    "print(\"Size of mixed 4-6-bit quantized model\")\n",
    "print(\"============================\")\n",
    "print(f\"{get_directory_size_mb(mlx_path):2.4f} GB\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0b63ef3",
   "metadata": {},
   "source": [
    "### Model fine-tuning\n",
    "Let's use the mistral-7b-v0.3-4bit model we just quantized who won the latest Super Bowl. As expected, the answer is correct but outdated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5efb794d",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mlx_lm.generate --model \"./mistral-7b-v0.3-4bit\" \\\n",
    "    --prompt \"Who played in the latest super bowl?\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85a72fd2",
   "metadata": {},
   "source": [
    "Let's train on a small dataset with questions and answers about the latest Super Bowl, we can update the model’s knowledge and have it answer accurately."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4c31126",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !mlx_lm.lora --model \"./mistral-7b-v0.3-4bit\" --train --data ./data --iters 300 --batch-size 8 --mask-prompt --learning-rate 1e-5\n",
    "\n",
    "if os.path.exists(\"./adapters\"):\n",
    "    print(\"Size of adapters\")\n",
    "    print(\"================\")\n",
    "    print(f\"{get_directory_size_mb(\"./adapters\")*1024:2.2f} MB\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc1a1e1f",
   "metadata": {},
   "source": [
    "We can now ask the model the same question and it will provide us with an answer using new knowledge from the adapter."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7dcf9874",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mlx_lm.generate --model \"./mistral-7b-v0.3-4bit\" \\\n",
    "                 --prompt \"Who played in the latest super bowl?\" \\\n",
    "                 --adapter \"adapters\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a66c78f",
   "metadata": {},
   "source": [
    "After the training is complete you can fuse the adapter into the model using the `mlx_lm.fuse` command."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8935f7b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mlx_lm.fuse --model \"./mistral-7b-v0.3-4bit\" \\\n",
    "            --adapter-path \"adapters\" \\\n",
    "            --save-path \"fused-mistral-7b-v0.3-4bit\" \\"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e1bedc2",
   "metadata": {},
   "source": [
    "And we can test the fused model again for knowledge it has learned from the fine-tuning process"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "343a8977",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mlx_lm.generate --model \"./fused-mistral-7b-v0.3-4bit\" \\\n",
    "                 --prompt \"Who played in the latest super bowl?\" \\\n",
    "                 --temp 0.6"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.17"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
